import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.stats import kendalltau


def score_stability(df):
    """
    :param df: 输入一个DataFrame格式的数据，包含每个学生在不同次考试中的成绩信息
    :return: 返回不同次考试中排名之间的Kendall Tau相关系数矩阵
    """

    # 计算每个学生在不同次考试中的排名
    rank_df = df.rank(axis=0, method='min', ascending=False)
    # print(rank_df)
    # 计算不同次考试中排名之间的相关系数矩阵
    n_exams = len(rank_df.columns)
    # print(n_exams, rank_df.columns)
    print(pd.DataFrame)
    corr_matrix = pd.DataFrame(index=range(1, n_exams), columns=range(1, n_exams))
    for i in range(1, n_exams):
        for j in range(i + 1, n_exams + 1):
            corr, _ = kendalltau(rank_df.iloc[:, i - 1], rank_df.iloc[:, j - 1])
            corr_matrix.loc[i, j - 1] = corr
            corr_matrix.loc[j - 1, i] = corr

    # 输出相关系数矩阵
    print('\n不同次考试中排名之间的Kendall Tau相关系数矩阵:')
    print(corr_matrix.round(2))
    corr_matrix = corr_matrix.astype(float)

    # 计算上三角区域的相关系数均值
    upper_triangle = corr_matrix.values[np.triu_indices_from(corr_matrix, k=1)]
    mean_corr = np.mean(upper_triangle)

    # 打印相关系数矩阵和稳定性评价
    print('\n不同次考试中排名之间的Kendall Tau相关系数矩阵:')
    print(corr_matrix.round(2))
    print(f'平均相关系数为 {mean_corr:.2f}')

    if mean_corr >= 0.5:
        print('该科目考试班级内部稳定性强')
    elif 0 < mean_corr < 0.5:
        print('该科目考试班级内部稳定性一般')
    else:
        print('该科目考试班级内部稳定性有待提高')
    # 可视化相关系数矩阵
    plt.rcParams['font.family'] = 'STSong'
    mask = np.zeros_like(corr_matrix)
    # mask[np.triu_indices_from(mask)] = True
    with sns.axes_style("white"):
        ax = sns.heatmap(corr_matrix, mask=mask, annot=True, cmap='coolwarm', vmin=-1, vmax=1, square=True)
        ax.set_title('kendall matrix')
        plt.xticks(rotation=45)
        plt.show()

    return corr_matrix


def main():
    df = pd.read_excel("../TestExample/test_score001.xlsx")
    score_stability(df)


if __name__ == "__main__":
    main()
